import os
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent
from langgraph.graph import StateGraph, START, END
from IPython.display import Image, display

base_url = os.environ.get("DEEPSEEK_BASE_URL")
model_name = os.environ.get("DEEPSEEK_MODEL_NAME")
os.environ["OPENAI_API_KEY"] = os.getenv("DEEPSEEK_API_KEY")

llm = ChatOpenAI(base_url=base_url, model_name=model_name)

# 使用StateGraph接收一个字典，代表state的类型是dict
builder = StateGraph(dict)


# print(f'builder：{builder}')
# print(f'builder.schemas：{builder.schemas}')

# 定义节点
def addition(state):
    print(f"state: {state}")
    return {"x": state["x"] + 1}


def subtraction(state):
    print(f"state: {state}")
    return {"x": state["x"] - 2}


builder.add_node("addition", addition)
builder.add_node("subtraction", subtraction)

builder.add_edge(START, "addition")
builder.add_edge("addition", "subtraction")
builder.add_edge("subtraction", END)

# print(f"edges: {builder.edges}")

# 通过builder创建图，并传入state
graph = builder.compile()

# display(Image(graph.get_graph(xray=True).draw_mermaid_png()))

# graph = create_react_agent(
#     model=llm,
#     tools=[],
#     prompt='你是一个智能AI助手，名字是小团团。你的性格活泼、开朗、乐观，是大家的开心果。请用符合你性格的语言回答问题。如果有不知道的问题，可以调用相关的tool辅助搜索答案。'
# )
